# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: inkyu

import json
import math
import logging

from tornado import gen
import requests

from coin.base.param_util import to_list
from coin.base.timestamp import get_timestamp
import coin.flow.subscriber as flow_subscriber
import coin.proto.coin_order_enums_pb2 as coin_order

from coin.exchange.base.order_gateway import (OrderType,
                                              BaseOrderGatewayInfoMutator,
                                              OrderGatewayBase,
                                              OrderGatewayStatus)
from coin.exchange.base.order_container import OrderContainer
from coin.exchange.base.order_id_util import OrderIdGenerator

from coin.exchange.base.order_gateway_logger import gen_og_log_request
from coin.exchange.upbit_v1.order_gateway_log import UpbitOrderGatewayLogProvider

from coin.exchange.kr_rest.product_holders import get_holder_from_product
from coin.exchange.upbit_v1.order import (create_upbit_order, create_upbit_order_from_snapshot)
from coin.exchange.upbit_v1.kr_rest.product import UpbitProduct
from coin.exchange.upbit_v1.kr_rest.currency import UpbitCurrency
from coin.exchange.upbit_v1.rest.private import UpbitPrivateQueryer
from coin.exchange.upbit_v1.rest.private_client import UpbitPrivateClient
from coin.exchange.upbit_v1.rest.private_parser import UpbitPrivateParser, UpbitExchangeOrderState
from coin.exchange.upbit_v1.kr_rest.private_client import (UpbitPrivateParser as
                                                           UpbitPrivateParserProto)
from coin.proto.coin_query_pb2 import AccountBalance
from coin.proto.coin_order_gateway_pb2 import OrderEvent
from coin.exchange.base.rest.task_scheduler import Scheduler
from coin.strategy.hk_hard_arb.og_util import print_working_orders


class UpbitGatewayInfo(BaseOrderGatewayInfoMutator):
  CurrencyType = UpbitCurrency
  ProductType = UpbitProduct
  BalanceType = float
  PositionType = None

  def __init__(self, currencies, products):
    super().__init__(currencies, products)

  def balance_as_proto(self):
    return AccountBalance(exchange='Upbit',
                          market_type="Spot",
                          each_balance=self.get_notNone_balances_internal())


class UpbitOrderGateway(OrderGatewayBase):
  _exchange = 'Upbit'
  _market_type = 'Spot'
  _api_version = None

  OGInfoType = UpbitGatewayInfo

  def __init__(self, currencies, products, config, logger=None):
    super().__init__()
    self._logger = logger or logging.getLogger(__name__)
    self._config = config

    self._og_info = UpbitGatewayInfo(to_list(currencies), to_list(products))
    self._orders = OrderContainer()

    self._order_formatters = {
        product: get_holder_from_product(product).get_formatter()
        for product in self._og_info.products
    }

    self._balance_ready = False
    self._orders_ready = False
    self._order_queries_done = {state: False for state in UpbitExchangeOrderState}
    self._order_id_generator = OrderIdGenerator()

    self._rest_client = UpbitPrivateClient(config, max_connections=6)

    # OG Logger
    self._og_logger = UpbitOrderGatewayLogProvider(gen_og_log_request(config, self))

    # Following instance attributes are initialized in start method.
    self._ioloop = None
    self._flow_sub = None
    self._private_queryer = None
    self._scheduler = None
    self._maker_time_threshold = 2

  def get_working_order(self):
    return self._orders.working_orders.values()

  def renew_order_formatter(self):
    self._order_formatters = {
        product: get_holder_from_product(product).get_formatter()
        for product in self._og_info.products
    }

  def has_pending_order(self, product):
    assert isinstance(product, UpbitProduct), product.__class__
    for order in self._orders.working_orders.values():
      if order.product == product and not order.accepted:
        return True
    return False

  def _check_working_orders(self):
    print_working_orders(self._logger, self)
    timestamp = get_timestamp()
    for order in self.get_working_order():
      if order.cancel_sent is True:
        diff_time = (timestamp - order.internal.cancel_sent_time) / 1e9
        if diff_time < 10:
          continue

        order.cancel_sent = False  # So this order can be cancelled again.
        order.internal.cancel_sent_count += 1
        order.internal.cancel_sent_time = None
        self._logger.warning("Reset order %s's cancel_sent flag. cancel_sent_count = %d",
                             order.order_id,
                             order.internal.cancel_sent_count)
      if order.internal.last_seen_time:
        last_seen_elapsed = (timestamp - order.internal.last_seen_time) / 1e9
        if last_seen_elapsed > 60:
          self._orders.remove_from_working_order(order.order_id)
          self._logger.error('Remove disappeared order!')

  def start(self, queue, ioloop, flow_sub=None):
    self._ioloop = ioloop
    self._private_queryer = UpbitPrivateQueryer(
        self._config, queue, '1', [product.native_symbol for product in self._og_info.products])
    self._private_queryer.start(ioloop)
    self._scheduler = Scheduler(logger=self._logger)
    self._scheduler.start()
    self._flow_sub = flow_sub or flow_subscriber.from_queue(queue)
    self._flow_sub.subscribe('upbit_rest_private_1', self.on_rest_private_msg)

  def _check_ready(self):
    if self._balance_ready and self._orders_ready:
      self._og_ready.set_status(OrderGatewayStatus.READY)
      job_name = 'working_orders_checker'
      if self.is_ready() and not self._scheduler.has_job_name(job_name):
        self._scheduler.add_periodic_task(
            period=10,
            job_name=job_name,
            job_func=self._check_working_orders,
        )
        self._logger.info('Add working order list periodic checker!')

  def _handle_rest_balance(self, message):
    self._og_info.set_zero_balances()
    account_balance = UpbitPrivateParserProto.parse_balance(message)
    for currency_balance in account_balance.each_balance:
      currency = UpbitCurrency.FromStrNativeCurrencyNothrow(currency_balance.currency_native)
      if currency is not None:
        self._og_info.set_balance(currency, currency_balance.total, currency_balance)

    self._og_logger.write_balance(self._og_info.balance_as_proto())
    self._balance_ready = True
    self._check_ready()

  def _handle_order_snapshot(self, order_snapshot):
    order = self._orders.by_exchange_order_id(order_snapshot.exchange_order_id)
    unknown = False
    timestamp = get_timestamp()
    if order:
      assert math.isclose(order.price, order_snapshot.price_orig), (order, order_snapshot)
      assert order.order_type == order_snapshot.order_type, (order, order_snapshot)
      assert order.order_side == order_snapshot.order_side, (order, order_snapshot)

      order.accepted = True
      order.internal.last_seen_time = timestamp
      if order_snapshot.state == UpbitExchangeOrderState.CANCEL:
        if order.internal.state == UpbitExchangeOrderState.CANCEL:
          pass
        elif order.internal.state == UpbitExchangeOrderState.WAIT:
          self._logger.info('Cancel confirmed: %s(%s)',
                            order.order_id,
                            order_snapshot.exchange_order_id)
          self._orders.remove_from_working_order(order.order_id)
          self._og_logger.gen_cancel_response_and_log(order.product.symbol, True, order.order_id)
          self._og_logger.gen_order_event_and_log(order, timestamp, OrderEvent.CANCEL_CONFIRMED)
        elif order.internal.state == UpbitExchangeOrderState.DONE:
          self._logger.error('Cancel after full execution: %s(%s)', order_snapshot)
        else:
          raise ValueError(order.internal.state)
      elif order_snapshot.state == UpbitExchangeOrderState.WAIT:
        pass
      elif order_snapshot.state == UpbitExchangeOrderState.DONE:
        if order.internal.state == UpbitExchangeOrderState.DONE:
          pass
        elif order.internal.state == UpbitExchangeOrderState.CANCEL:
          self._logger.error('Fill after cancelation: %s', order_snapshot)
        elif order.internal.state == UpbitExchangeOrderState.WAIT:
          self._orders.remove_from_working_order(order.order_id)
        else:
          raise ValueError(order.internal.state)
      else:
        raise ValueError(order_snapshot)
      order.internal.state = order_snapshot.state
    else:
      order_id = self._order_id_generator.gen_order_id(is_foreign=True)
      product = UpbitProduct.FromStrNativeProductNothrow(order_snapshot.native_symbol)
      # Ignore irrelevant product
      if product is None or product not in self._og_info.products:
        return

      order = create_upbit_order_from_snapshot(order_snapshot=order_snapshot,
                                               order_id=order_id,
                                               timestamp=timestamp)
      unknown = True
      working = False
      if order_snapshot.state == UpbitExchangeOrderState.CANCEL:
        if self._orders_ready:
          self._logger.error('Unknown cancellation: %s', order_snapshot)
      elif order_snapshot.state == UpbitExchangeOrderState.WAIT:
        working = True
        self._logger.info('Accepted foreign order: %s', order_id)
      elif order_snapshot.state == UpbitExchangeOrderState.DONE:
        pass
      else:
        raise ValueError(order_snapshot)
      self._orders.add(order, working=working)

    order.internal.timestamp = order_snapshot.timestamp_created
    if order.tag is None:
      order.tag = {'submit_ack_timestamp': timestamp}

    if order_snapshot.qty_filled > 0.:
      price_avg_fill = order_snapshot.paid_fee / 5.e-4 / order_snapshot.qty_filled
    else:
      price_avg_fill = order_snapshot.price_orig
    last_fill = order.internal.fill.update_by_accum_fill(price_avg_fill=price_avg_fill,
                                                         qty_accum_fill=order_snapshot.qty_filled)
    if self._orders_ready and last_fill is not None:
      if unknown:
        self._logger.info('Fill from unknown order: %s', order.order_id)
      self._logger.info('Upbit FILL (%s, %s): %s %f,%f/%f @ %f(%f)',
                        order.product,
                        order.order_id,
                        order.order_side,
                        last_fill.qty_last_fill,
                        order_snapshot.qty_filled,
                        order.qty,
                        last_fill.price_last_fill,
                        order.price)

      fill_type = coin_order.TAKER_FILL_TYPE
      if math.isclose(last_fill.price_last_fill, order.price) and order.tag:
        order_lifetime = (timestamp - order.tag['submit_ack_timestamp']) / 1e9
        if order_lifetime > self._maker_time_threshold:
          fill_type = coin_order.MAKER_FILL_TYPE
      self._publish_fill(last_fill.price_last_fill, last_fill.qty_last_fill, order)
      self._og_logger.gen_order_event_and_log(order,
                                              timestamp,
                                              OrderEvent.ORDER_FILLED,
                                              filled_qty=last_fill.qty_last_fill,
                                              guess_fill_type=fill_type)

  def _handle_rest_open_orders(self, message, *, state):
    order_snapshots = UpbitPrivateParser.parse_orders(message)
    exchange_order_id_list = set()
    for order_snapshot in order_snapshots:
      exchange_order_id_list.add(order_snapshot.exchange_order_id)
      self._handle_order_snapshot(order_snapshot)

    for order in self.get_working_order():
      if not order.accepted:
        continue

      if order.internal.exchange_order_id not in exchange_order_id_list:
        if order.internal.cancel_sent_count > 5:
          self._orders.remove_from_working_order(order.order_id)
          self._logger.error(
              'Remove cancel failed order (%s, %s), '
              'which is not in the exchange working order list anymore!',
              order.order_id,
              order.internal.exchange_order_id)

    if not self._orders_ready:
      self._order_queries_done[state] = True
      if all(self._order_queries_done.values()):
        self._orders_ready = True
        self._check_ready()

  def on_rest_private_msg(self, record, queue_data, topic_data):
    message = record.data

    query = message['query']
    if query == 'accounts':
      self._handle_rest_balance(message['data'])
    elif query == 'orders':
      self._handle_rest_open_orders(message['data'], state=message['state'])
    else:
      raise ValueError('Unknown query: %s' % query)

  def submittable(self):
    return self.is_ready() and self._rest_client.rate_limiter.order_available()

  def submit_impl(self, product, price, qty, order_side, order_type, order_id=None, tag=None):
    assert self.is_ready()
    assert order_type == OrderType.LIMIT
    assert order_id is None  # Not supported yet
    assert tag is None

    order_id = order_id or self._order_id_generator.gen_order_id()
    order = create_upbit_order(order_id=order_id,
                               product=product,
                               price=price,
                               qty=qty,
                               order_side=order_side,
                               order_type=order_type,
                               tag=tag)
    self._orders.add(order)
    self._ioloop.add_callback(self._submit_impl, order)
    timestamp = get_timestamp()
    self._og_logger.gen_order_request_and_order_event_and_log(product.symbol,
                                                              price,
                                                              qty,
                                                              order_side,
                                                              order_type,
                                                              order_id,
                                                              timestamp)

  @gen.coroutine
  def _submit_impl(self, order):
    message = None
    response = None
    try:
      formatter = self._order_formatters[order.product]
      price_str = formatter.format_price(order.price)
      qty_str = formatter.format_qty(order.qty)
      response = yield self._rest_client.submit_order(native_symbol=order.product.native_symbol,
                                                      order_side=order.order_side,
                                                      price_str=price_str,
                                                      qty_str=qty_str,
                                                      order_type=order.order_type)
      message = json.loads(response.content)
    except json.JSONDecodeError:
      self._logger.exception('JsonDecodeException: %s', response.content)
    except requests.exceptions.RequestException:
      self._logger.exception('RequestException')
    except Exception as e:
      self._logger.exception('Unknown Exception, %s', e)

    timestamp = get_timestamp()
    if message:
      if response.status_code == 201:  # CREATED
        order_snapshot = UpbitPrivateParser.parse_order(message)
        order_foreign = self._orders.pop_by_exchange_order_id(order_snapshot.exchange_order_id)
        self._orders.assign_exchange_order_id(order, order_snapshot.exchange_order_id)
        if order_foreign is not None:
          self._orders.remove_from_working_order(order_foreign.order_id)
          order.internal.fill = order_foreign.internal.fill
          self._logger.info('Replace foreign order %s -> %s',
                            order_foreign.order_id,
                            order.order_id)
        self._handle_order_snapshot(order_snapshot)

        self._og_logger.gen_order_response_and_order_event_and_log(order.product.symbol,
                                                                   True,
                                                                   order.order_id,
                                                                   timestamp)
        self._og_logger.gen_order_event_and_log(order, timestamp, OrderEvent.ORDER_ACCEPTED)
        return
      else:
        self._logger.error('Invalid status code: %s (%s)', response.status_code, response.content)

    self._orders.remove_from_working_order(order.order_id)
    self._logger.error('Submit failed: %s', order.order_id)
    self._og_logger.gen_order_response_and_order_event_and_log(order.product.symbol,
                                                               False,
                                                               order.order_id,
                                                               timestamp)

  @gen.coroutine
  def _cancel_by_exchange_order_id(self, exchange_order_id):
    message = None
    response = None
    try:
      response = yield self._rest_client.cancel_order(exchange_order_id)
      message = json.loads(response.content)
    except requests.exceptions.RequestException:
      self._logger.exception('RequestException')
    except json.JSONDecodeError:
      self._logger.exception('JsonDecodeError')

    order = self._orders.by_exchange_order_id(exchange_order_id)
    if response and response.status_code == 200:  # OK
      order.cancel_sent = True
      order_snapshot = UpbitPrivateParser.parse_order(message)
      if order_snapshot.state != UpbitExchangeOrderState.CANCEL:
        self._logger.debug('Cancel successful, order state is not `CANCEL`.\n%s', message)
        if order.internal.cancel_sent_count > 2:
          self._logger.warning('Cancelled multiple times: %s', message)

      self._logger.info('Cancel sent: %s(%s)', order.order_id, exchange_order_id)
      self._handle_order_snapshot(order_snapshot)
    else:
      error = message.get('error', None)
      if error:
        # Reset cancel_sent flag.
        order.cancel_sent = False
        order.internal.cancel_sent_time = None
        error_name = error.get('name', None)
        # Recent upbit message uses 'order_not_found' rather than
        # 'V1::Exceptions::OrderNotFound'
        if error_name in ('V1::Exceptions::OrderNotFound', 'order_not_found'):
          self._orders.remove_from_working_order(order.order_id)
          self._logger.warning('Order not found, deleting: %s', order.order_id)
          return
        else:
          self._logger.error('Cancel failed, error_name: %s', error_name)
      self._logger.warning('Cancellation failed: %s %s', exchange_order_id, message)

  def cancel_impl(self, order_id, *, ignore_error=False):
    order = self._orders.get_working_order(order_id)
    exchange_order_id = order.internal.exchange_order_id
    timestamp = get_timestamp()
    if order.cancel_sent:
      self._logger.debug('Already cancelled %s' % order_id)
      return
    if exchange_order_id:
      order.cancel_sent = True
      order.internal.cancel_sent_time = timestamp
      self._ioloop.add_callback(self._cancel_by_exchange_order_id, exchange_order_id)
      self._og_logger.gen_cancel_request_and_order_event_and_log(order.product.symbol,
                                                                 order.order_id,
                                                                 timestamp)
    else:
      self._logger.error('Order is not submitted: %s', order_id)

  def cancel_product_impl(self, product, *, ignore_error):
    order_ids = []
    self._logger.debug('Canceling %s', product)
    for order in self._orders.working_orders.values():
      if order.product == product:
        order_ids.append(order.order_id)
    self.cancel_multiple(order_ids, ignore_error=ignore_error)

  def cancel_all_impl(self, *, ignore_error=False):
    order_ids = []
    for order in self._orders.working_orders.values():
      order_ids.append(order.order_id)
    for order_id in order_ids:
      self.cancel(order_id)

  def cancel_multiple_impl(self, order_ids, *, ignore_error=False):
    for order_id in order_ids:
      self.cancel(order_id)
